import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd

def plot_energy_trends(df, energy_sources):
    """
    绘制能源趋势图
    
    参数:
        df: 数据DataFrame
        energy_sources: 能源列表
    """
    fig, axes = plt.subplots(nrows=len(energy_sources), ncols=1, figsize=(15, 15), sharex=True)
    
    for i, source in enumerate(energy_sources):
        axes[i].plot(df['DateTime'], df[source])
        axes[i].set_title(f'{source} Trend')
        axes[i].set_ylabel('MW')
        axes[i].grid(True)
    
    plt.tight_layout()
    plt.show()

def plot_energy_distribution(df, energy_sources):
    """
    绘制能源分布图
    
    参数:
        df: 数据DataFrame
        energy_sources: 能源列表
    """
    plt.figure(figsize=(12, 8))
    
    # 创建箱线图
    sns.boxplot(data=df[energy_sources])
    plt.title('Energy Sources Distribution')
    plt.ylabel('MW')
    plt.xticks(rotation=45)
    plt.grid(True, axis='y')
    
    plt.tight_layout()
    plt.show()

def plot_hourly_consumption_production(df):
    """
    绘制每小时用电量和发电量的线图
    
    参数:
        df: 数据DataFrame
    """
    plt.figure(figsize=(15, 6))
    
    plt.plot(df['DateTime'], df['Consumption'], label='Consumption')
    plt.plot(df['DateTime'], df['Production'], label='Production')
    
    plt.title('Hourly Electricity Consumption and Production')
    plt.xlabel('Date')
    plt.ylabel('MW')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

def plot_stacked_area(df, energy_sources):
    """
    绘制堆叠面积图显示每小时的电力生产来源
    
    参数:
        df: 数据DataFrame
        energy_sources: 能源列表
    """
    plt.figure(figsize=(15, 8))
    
    plt.stackplot(df['DateTime'], 
                 [df[source] for source in energy_sources],
                 labels=energy_sources,
                 alpha=0.8)
    
    plt.title('Hourly Electricity Production by Source')
    plt.xlabel('Date')
    plt.ylabel('MW')
    plt.legend(loc='upper left')
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

def plot_lag_features(df, column, lags=24):
    """
    绘制滞后特征图
    
    参数:
        df: 数据DataFrame
        column: 要分析的列
        lags: 滞后数量
    """
    plt.figure(figsize=(12, 8))
    
    for i in range(1, lags+1):
        plt.scatter(df[column].shift(i), df[column], alpha=0.5, s=5, label=f'Lag {i}')
    
    plt.title(f'Lag Features for {column}')
    plt.xlabel(f'Lagged {column}')
    plt.ylabel(column)
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

def plot_correlation_matrix(df, columns):
    """
    绘制相关矩阵
    
    参数:
        df: 数据DataFrame
        columns: 要包含在相关矩阵中的列
    """
    plt.figure(figsize=(12, 10))
    
    corr = df[columns].corr()
    mask = np.triu(np.ones_like(corr, dtype=bool))
    
    sns.heatmap(corr, mask=mask, annot=True, cmap='coolwarm', vmin=-1, vmax=1, center=0,
                square=True, linewidths=.5, cbar_kws={"shrink": .5})
    
    plt.title('Correlation Matrix')
    plt.tight_layout()
    plt.show()

def plot_model_predictions(actual, predictions, title, scaler=None):
    """
    绘制模型预测结果
    
    参数:
        actual: 实际值
        predictions: 预测值
        title: 图表标题
        scaler: 如果提供，用于反向转换缩放的数据
    """
    plt.figure(figsize=(15, 6))
    
    if scaler is not None:
        actual = scaler.inverse_transform(actual)
        predictions = scaler.inverse_transform(predictions)
    
    plt.plot(actual, label='Actual')
    plt.plot(predictions, label='Predicted', alpha=0.7)
    
    plt.title(title)
    plt.xlabel('Time')
    plt.ylabel('MW')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()